/*
Copyright (c) 2011, Nimar S. Arora
All rights reserved.
*/
#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#include "network.h"
#include "statistics.h"

static int sample_weights(int num_wts, double * p_wts)
{
  int i;
  double total;
  double target;
  
  total = 0;
  
  for (i=0; i<num_wts; i++)
    total += p_wts[i];

  if (0 == total)
    return rand() % num_wts;
  
  target = total * rand() / RAND_MAX;
  
  total = 0;
  
  for (i=0; i<num_wts; i++)
  {
    total += p_wts[i];
    if (total > target)
      return i;
  }

  return num_wts-1;
}

static int sample_log_weights(int num_wts, double * p_log_wts)
{
  int i;
  double max_log_wt;
  
  /* compute the maximum log weight */
  max_log_wt = -HUGE_VAL;
  for (i=0; i<num_wts; i++)
    if (p_log_wts[i] > max_log_wt)
      max_log_wt = p_log_wts[i];

  if (isinf(max_log_wt))
    return rand() % num_wts;
  
  /* rescale the weights and exponentiate */
  for (i=0; i<num_wts; i++)
    p_log_wts[i] = exp(p_log_wts[i] - max_log_wt);
  
  /* now, sample from these regular weights */
  return sample_weights(num_wts, p_log_wts);
}

static void rand_init(Graph_t * p_graph, Instance_t * p_inst)
{
  int v;
  
  for (v=0; v<p_graph->graph_num_vars; v++)
  {
    Variable_t * p_var;
    
    if (p_inst->inst_p_is_obs[v])
      continue;
      
    p_var = p_graph->graph_p_vars + v;
      
    p_inst->inst_p_val[v] = rand() % p_var->var_card;
  }
}
  
/* go through each variable and sample a value for it using Gibbs */
static void kernel_1var(Graph_t * p_graph, Instance_t * p_inst,
                        double * p_weights)
{
  int vi;                                    /* variable index */

  /* scan each non-observed variable */
  /* IMPROVEME: keep a list of non-observed variables */
  for (vi=0; vi<p_graph->graph_num_vars; vi++)
  {
    Variable_t * p_var;
    int value;
        
    if (p_inst->inst_p_is_obs[vi])
      continue;
        
    p_var = p_graph->graph_p_vars + vi;
        
    /* compute a weight for each possible value of the variable */
    for (value=0; value<p_var->var_card; value ++)
    {
      int fi;                            /* factor index */
          
      p_inst->inst_p_val[vi] = value;
      /* compute the weight for this value */
      p_weights[value] = 0;

      for (fi=0; fi<p_var->var_num_facs; fi++)
      {
        Factor_t * p_fac;
            
        p_fac = p_var->var_p_facs[fi];

        p_weights[value] += graph_factor(p_fac, p_inst->inst_p_val);
      }
    }

    /* randomly choose a value proportional to the weights */
    value = sample_log_weights(p_var->var_card, p_weights);

    p_inst->inst_p_val[vi] = value;
  }
}

void mcmc(Graph_t * p_graph, Instance_t * p_inst, int numsamples, int lag, 
          int anneal)
{
  int samplenum;
  double * p_weights;
  Stat_PE_t stat;

  /* randomly initialize the starting state */
  rand_init(p_graph, p_inst);
  
  /* initialize memory for MCMC weights */
  p_weights = (double *) malloc(MAX_CARD * sizeof(*p_weights));

  stat_PE_init(&stat, p_graph);
  
  for (samplenum=0; samplenum < (numsamples * lag); samplenum ++)
  {
    kernel_1var(p_graph, p_inst, p_weights);

    if (0 == ((samplenum+1) % lag))
    {
      stat_PE_update(&stat, p_graph, p_inst);

      /* we will print 10 incremental values for easy plotting */
      if (((samplenum+1) % (numsamples * lag / 10)) == 0)
      {
        printf("samplenum %d: ", (samplenum+1)/lag);
        stat_PE_write(&stat, p_graph, stdout);
      }
    }
  }
  stat_PE_write(&stat, p_graph, stdout);
}

